'''
X = “在 n 个人中，满足三人同一天的 unordered triplets 的数量”（每个三人组合若三人同天则计 1）

lambda = E[X] = C(n, 3)/365^2
C(n,3) = n*(n-1)*(n-2)/6

P(X=k) = lambda^k * e^(-lambda) / k!

P(X=1) = lambda * e^(-lambda)

'''

import math
import numpy as np
import matplotlib.pyplot as plt
import numpy.random as npr
import matplotlib
matplotlib.use(backend="TkAgg")

def poisson_triplet_prob(n, m=365, k=1):
    if n < 3:
        return 0
    lam = math.comb(n,3)/m**2
    return math.exp(-lam)*(lam**k)/math.factorial(k)

'''
步骤2：统计每天生日人数
counts = np.array([np.bincount(row, minlength=365) for row in B])
对每行统计各生日出现的次数。
第一行的统计结果：
生日123：3人（位置0,2,4）
生日45：2人（位置1,6）
生日78：2人（位置3,8）
生日200：1人（位置5）
生日300：1人（位置7）
生日150：1人（位置9）
其他362天：0人

步骤3：计算三人组合数
triplet_counts = (counts * (counts - 1) * (counts - 2) // 6).sum(axis=1)
对每个生日人数k，计算C(k,3) = k(k-1)(k-2)/6
第一行的计算：
生日123：C(3,3) = 3×2×1/6 = 1
生日45：C(2,3) = 2×1×0/6 = 0
生日78：C(2,3) = 0
其他：都是0
总三人组合数 = 1
'''
# monte carlo estimate for P(X=1)
def monte_carlo_p_x1(n, trials=5000,m=365,seed=42):
    rng = npr.default_rng(seed)
    B=rng.integers(low=0, high=m, size=(trials, n))
    counts = np.array([np.bincount(row, minlength=m) for row in B], dtype=np.int16)
    triplet_counts = (counts*(counts-1)*(counts-2)//6).sum(axis=1)
    return (triplet_counts==1).mean()

# Generate curves
N_values = range(10, 201)
poisson_values = [poisson_triplet_prob(n) for n in N_values]
mc_values = [monte_carlo_p_x1(n, trials=5000) for n in N_values]  # smaller trials for speed

plt.plot(list(N_values), poisson_values, label="Poisson近似 P(X==1)")
plt.plot(list(N_values), mc_values, label="Monte Carlo估计 P(X==1)")
plt.xlabel("人数 N")
plt.ylabel("恰好一个三人组合的概率")
plt.title("P(X==1) 随人数变化 (n=1~200)")
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()